# 数据库管理模块

import os
import sys
# 将项目根目录添加到Python搜索路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import sqlite3
import pymongo
from datetime import datetime
from config import settings
from utils import setup_logger, create_dir_if_not_exists

logger = setup_logger('db_manager')


class DBManager:
    """数据库管理类，支持SQLite和MongoDB"""
    def __init__(self):
        self.db_type = settings.DB_TYPE
        self.db = None
        self.connect()
        if self.db_type == 'sqlite':
            self.create_tables()

    def connect(self):
        """连接数据库"""
        try:
            if self.db_type == 'sqlite':
                # 确保SQLite数据库目录存在
                db_dir = os.path.dirname(settings.SQLITE_DB_PATH)
                create_dir_if_not_exists(db_dir)
                # 连接SQLite数据库
                self.db = sqlite3.connect(settings.SQLITE_DB_PATH)
                logger.info(f"成功连接SQLite数据库: {settings.SQLITE_DB_PATH}")
            elif self.db_type == 'mongodb':
                # 连接MongoDB数据库
                client = pymongo.MongoClient(settings.MONGODB_URI)
                self.db = client[settings.MONGODB_DB_NAME]
                # 测试连接
                client.server_info()
                logger.info(f"成功连接MongoDB数据库: {settings.MONGODB_DB_NAME}")
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            self.db = None

    def create_tables(self):
        """创建SQLite数据表"""
        if self.db_type != 'sqlite' or self.db is None:
            return

        try:
            cursor = self.db.cursor()
            # 创建歌曲表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS songs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    song_id TEXT UNIQUE,
                    title TEXT NOT NULL,
                    artist TEXT,
                    album TEXT,
                    duration TEXT,
                    file_path TEXT,
                    file_size INTEGER,
                    download_time TIMESTAMP,
                    status TEXT DEFAULT 'pending',
                    url TEXT,
                    lyrics TEXT,
                    cover_url TEXT,
                    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 创建视频表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS videos (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    video_id TEXT UNIQUE,
                    title TEXT NOT NULL,
                    artist TEXT,
                    duration TEXT,
                    file_path TEXT,
                    file_size INTEGER,
                    download_time TIMESTAMP,
                    status TEXT DEFAULT 'pending',
                    url TEXT,
                    cover_url TEXT,
                    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    media_type TEXT DEFAULT 'video'
                )
            ''')
            
            # 创建分类表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS categories (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT UNIQUE,
                    url TEXT,
                    description TEXT,
                    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 创建歌曲分类关联表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS song_categories (
                    song_id INTEGER,
                    category_id INTEGER,
                    PRIMARY KEY (song_id, category_id),
                    FOREIGN KEY (song_id) REFERENCES songs (id),
                    FOREIGN KEY (category_id) REFERENCES categories (id)
                )
            ''')
            
            # 创建视频分类关联表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS video_categories (
                    video_id INTEGER,
                    category_id INTEGER,
                    PRIMARY KEY (video_id, category_id),
                    FOREIGN KEY (video_id) REFERENCES videos (id),
                    FOREIGN KEY (category_id) REFERENCES categories (id)
                )
            ''')
            
            # 创建爬取记录表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS crawl_records (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    url TEXT,
                    status TEXT,
                    start_time TIMESTAMP,
                    end_time TIMESTAMP,
                    items_count INTEGER,
                    error_message TEXT
                )
            ''')
            
            self.db.commit()
            logger.info("SQLite数据表创建成功")
        except Exception as e:
            logger.error(f"创建数据表失败: {e}")
            self.db.rollback()

    def insert_song(self, song_data):
        """插入歌曲信息"""
        if self.db is None:
            return False

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                # 检查歌曲是否已存在
                cursor.execute("SELECT id FROM songs WHERE song_id = ?", (song_data.get('song_id'),))
                existing = cursor.fetchone()
                
                if existing:
                    # 更新现有歌曲
                    update_data = {
                        'title': song_data.get('title'),
                        'artist': song_data.get('artist'),
                        'album': song_data.get('album'),
                        'duration': song_data.get('duration'),
                        'file_path': song_data.get('file_path'),
                        'file_size': song_data.get('file_size'),
                        'download_time': song_data.get('download_time'),
                        'status': song_data.get('status'),
                        'url': song_data.get('url'),
                        'lyrics': song_data.get('lyrics'),
                        'cover_url': song_data.get('cover_url'),
                        'update_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    }
                    
                    # 构建更新语句
                    update_fields = ', '.join([f"{k} = ?" for k in update_data.keys()])
                    update_values = tuple(update_data.values()) + (song_data.get('song_id'),)
                    
                    cursor.execute(f"UPDATE songs SET {update_fields} WHERE song_id = ?", update_values)
                    song_id = existing[0]
                else:
                    # 插入新歌曲
                    insert_data = {
                        'song_id': song_data.get('song_id'),
                        'title': song_data.get('title'),
                        'artist': song_data.get('artist'),
                        'album': song_data.get('album'),
                        'duration': song_data.get('duration'),
                        'file_path': song_data.get('file_path'),
                        'file_size': song_data.get('file_size'),
                        'download_time': song_data.get('download_time'),
                        'status': song_data.get('status', 'pending'),
                        'url': song_data.get('url'),
                        'lyrics': song_data.get('lyrics'),
                        'cover_url': song_data.get('cover_url'),
                        'create_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                        'update_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    }
                    
                    # 构建插入语句
                    fields = ', '.join(insert_data.keys())
                    placeholders = ', '.join(['?' for _ in insert_data.keys()])
                    values = tuple(insert_data.values())
                    
                    cursor.execute(f"INSERT INTO songs ({fields}) VALUES ({placeholders})", values)
                    song_id = cursor.lastrowid
                
                self.db.commit()
                return song_id
            elif self.db_type == 'mongodb':
                # MongoDB插入或更新
                songs_collection = self.db['songs']
                
                # 准备数据
                mongo_data = song_data.copy()
                mongo_data['update_time'] = datetime.now()
                
                if not mongo_data.get('create_time'):
                    mongo_data['create_time'] = datetime.now()
                
                # 更新或插入
                result = songs_collection.update_one(
                    {'song_id': song_data.get('song_id')},
                    {'$set': mongo_data},
                    upsert=True
                )
                
                if result.upserted_id:
                    return str(result.upserted_id)
                else:
                    # 查询更新后的文档ID
                    song = songs_collection.find_one({'song_id': song_data.get('song_id')})
                    return str(song['_id']) if song else False
        except Exception as e:
            logger.error(f"插入/更新歌曲失败: {e}")
            if self.db_type == 'sqlite':
                self.db.rollback()
            return False

    def get_song_by_id(self, song_id):
        """根据ID获取歌曲信息"""
        if self.db is None:
            return None

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                cursor.execute("SELECT * FROM songs WHERE song_id = ?", (song_id,))
                song = cursor.fetchone()
                if song:
                    # 转换为字典
                    columns = [desc[0] for desc in cursor.description]
                    return dict(zip(columns, song))
            elif self.db_type == 'mongodb':
                songs_collection = self.db['songs']
                song = songs_collection.find_one({'song_id': song_id})
                if song:
                    # 转换ObjectId为字符串
                    song['_id'] = str(song['_id'])
                return song
        except Exception as e:
            logger.error(f"获取歌曲信息失败: {e}")
        return None

    def get_pending_songs(self, limit=100):
        """获取待下载的歌曲"""
        if self.db is None:
            return []

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                cursor.execute("SELECT * FROM songs WHERE status = 'pending' LIMIT ?", (limit,))
                songs = cursor.fetchall()
                # 转换为字典列表
                columns = [desc[0] for desc in cursor.description]
                return [dict(zip(columns, song)) for song in songs]
            elif self.db_type == 'mongodb':
                songs_collection = self.db['songs']
                songs = list(songs_collection.find({'status': 'pending'}).limit(limit))
                # 转换ObjectId为字符串
                for song in songs:
                    song['_id'] = str(song['_id'])
                return songs
        except Exception as e:
            logger.error(f"获取待下载歌曲失败: {e}")
        return []

    def check_song_exists(self, song_id):
        """检查歌曲是否已存在"""
        if self.db is None:
            return False

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                cursor.execute("SELECT id FROM songs WHERE song_id = ?", (song_id,))
                return cursor.fetchone() is not None
            elif self.db_type == 'mongodb':
                songs_collection = self.db['songs']
                return songs_collection.find_one({'song_id': song_id}) is not None
        except Exception as e:
            logger.error(f"检查歌曲是否存在失败: {e}")
        return False
        
    def update_song(self, song_id, song_data):
        """更新歌曲完整信息"""
        if self.db is None:
            return False

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                
                # 准备更新字段和值
                update_fields = []
                update_values = []
                
                # 遍历song_data中的每个键值对
                for key, value in song_data.items():
                    if key != 'song_id' and key != 'id' and key != 'create_time':
                        update_fields.append(f"{key} = ?")
                        update_values.append(value)
                
                # 添加更新时间
                update_fields.append("update_time = ?")
                update_values.append(datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
                
                # 添加song_id作为WHERE条件
                update_values.append(song_id)
                
                # 执行更新
                cursor.execute(f"UPDATE songs SET {', '.join(update_fields)} WHERE song_id = ?", update_values)
                self.db.commit()
                return cursor.rowcount > 0
            elif self.db_type == 'mongodb':
                songs_collection = self.db['songs']
                
                # 准备更新数据
                update_data = {'$set': song_data.copy()}
                update_data['$set']['update_time'] = datetime.now()
                
                # 移除不需要更新的字段
                for key in ['song_id', '_id', 'create_time']:
                    if key in update_data['$set']:
                        del update_data['$set'][key]
                
                # 执行更新
                result = songs_collection.update_one({'song_id': song_id}, update_data)
                return result.modified_count > 0
        except Exception as e:
            logger.error(f"更新歌曲信息失败: {e}")
            if self.db_type == 'sqlite':
                self.db.rollback()
        return False
        
    def update_video(self, video_id, video_data):
        """更新视频完整信息"""
        if self.db is None:
            return False

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                
                # 准备更新字段和值
                update_fields = []
                update_values = []
                
                # 遍历video_data中的每个键值对
                for key, value in video_data.items():
                    if key != 'video_id' and key != 'id' and key != 'create_time':
                        update_fields.append(f"{key} = ?")
                        update_values.append(value)
                
                # 添加更新时间
                update_fields.append("update_time = ?")
                update_values.append(datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
                
                # 添加video_id作为WHERE条件
                update_values.append(video_id)
                
                # 执行更新
                cursor.execute(f"UPDATE videos SET {', '.join(update_fields)} WHERE video_id = ?", update_values)
                self.db.commit()
                return cursor.rowcount > 0
            elif self.db_type == 'mongodb':
                videos_collection = self.db['videos']
                
                # 准备更新数据
                update_data = {'$set': video_data.copy()}
                update_data['$set']['update_time'] = datetime.now()
                
                # 移除不需要更新的字段
                for key in ['video_id', '_id', 'create_time']:
                    if key in update_data['$set']:
                        del update_data['$set'][key]
                
                # 执行更新
                result = videos_collection.update_one({'video_id': video_id}, update_data)
                return result.modified_count > 0
        except Exception as e:
            logger.error(f"更新视频信息失败: {e}")
            if self.db_type == 'sqlite':
                self.db.rollback()
        return False
        
    def update_song_status(self, song_id, status, file_path=None, file_size=None, error_message=None):
        """更新歌曲状态"""
        if self.db is None:
            return False

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                update_fields = ["status = ?", "update_time = ?"]
                update_values = [status, datetime.now().strftime('%Y-%m-%d %H:%M:%S')]
                
                if file_path:
                    update_fields.append("file_path = ?")
                    update_values.append(file_path)
                if file_size:
                    update_fields.append("file_size = ?")
                    update_values.append(file_size)
                if error_message:
                    # 如果数据库中有error_message字段，则更新它
                    # 检查表结构是否有error_message字段
                    cursor.execute("PRAGMA table_info(songs)")
                    columns = [col[1] for col in cursor.fetchall()]
                    if 'error_message' in columns:
                        update_fields.append("error_message = ?")
                        update_values.append(error_message)
                if status == 'downloaded':
                    update_fields.append("download_time = ?")
                    update_values.append(datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
                
                update_values.append(song_id)
                
                cursor.execute(f"UPDATE songs SET {', '.join(update_fields)} WHERE song_id = ?", update_values)
                self.db.commit()
                return cursor.rowcount > 0
            elif self.db_type == 'mongodb':
                songs_collection = self.db['songs']
                update_data = {
                    '$set': {
                        'status': status,
                        'update_time': datetime.now()
                    }
                }
                
                if file_path:
                    update_data['$set']['file_path'] = file_path
                if file_size:
                    update_data['$set']['file_size'] = file_size
                if error_message:
                    update_data['$set']['error_message'] = error_message
                if status == 'downloaded':
                    update_data['$set']['download_time'] = datetime.now()
                
                result = songs_collection.update_one({'song_id': song_id}, update_data)
                return result.modified_count > 0
        except Exception as e:
            logger.error(f"更新歌曲状态失败: {e}")
            if self.db_type == 'sqlite':
                self.db.rollback()
        return False

    def insert_crawl_record(self, url, status, start_time, end_time=None, items_count=0, error_message=None):
        """插入爬取记录"""
        if self.db is None:
            return False

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                cursor.execute(
                    "INSERT INTO crawl_records (url, status, start_time, end_time, items_count, error_message) VALUES (?, ?, ?, ?, ?, ?)",
                    (url, status, start_time, end_time, items_count, error_message)
                )
                self.db.commit()
                return cursor.lastrowid
            elif self.db_type == 'mongodb':
                crawl_records = self.db['crawl_records']
                record = {
                    'url': url,
                    'status': status,
                    'start_time': start_time,
                    'end_time': end_time,
                    'items_count': items_count,
                    'error_message': error_message
                }
                result = crawl_records.insert_one(record)
                return str(result.inserted_id)
        except Exception as e:
            logger.error(f"插入爬取记录失败: {e}")
            if self.db_type == 'sqlite':
                self.db.rollback()
        return False

    def insert_video(self, video_data):
        """插入视频信息"""
        if self.db is None:
            return False

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                # 检查视频是否已存在
                cursor.execute("SELECT id FROM videos WHERE video_id = ?", (video_data.get('video_id'),))
                existing = cursor.fetchone()
                
                if existing:
                    # 更新现有视频
                    update_data = {
                        'title': video_data.get('title'),
                        'artist': video_data.get('artist'),
                        'duration': video_data.get('duration'),
                        'file_path': video_data.get('file_path'),
                        'file_size': video_data.get('file_size'),
                        'download_time': video_data.get('download_time'),
                        'status': video_data.get('status'),
                        'url': video_data.get('url'),
                        'cover_url': video_data.get('cover_url'),
                        'media_type': video_data.get('media_type', 'video'),
                        'update_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    }
                    
                    # 构建更新语句
                    update_fields = ', '.join([f"{k} = ?" for k in update_data.keys()])
                    update_values = tuple(update_data.values()) + (video_data.get('video_id'),)
                    
                    cursor.execute(f"UPDATE videos SET {update_fields} WHERE video_id = ?", update_values)
                    video_id = existing[0]
                else:
                    # 插入新视频
                    insert_data = {
                        'video_id': video_data.get('video_id'),
                        'title': video_data.get('title'),
                        'artist': video_data.get('artist'),
                        'duration': video_data.get('duration'),
                        'file_path': video_data.get('file_path'),
                        'file_size': video_data.get('file_size'),
                        'download_time': video_data.get('download_time'),
                        'status': video_data.get('status', 'pending'),
                        'url': video_data.get('url'),
                        'cover_url': video_data.get('cover_url'),
                        'media_type': video_data.get('media_type', 'video'),
                        'create_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                        'update_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    }
                    
                    # 构建插入语句
                    fields = ', '.join(insert_data.keys())
                    placeholders = ', '.join(['?' for _ in insert_data.keys()])
                    values = tuple(insert_data.values())
                    
                    cursor.execute(f"INSERT INTO videos ({fields}) VALUES ({placeholders})", values)
                    video_id = cursor.lastrowid
                
                self.db.commit()
                return video_id
            elif self.db_type == 'mongodb':
                # MongoDB插入或更新
                videos_collection = self.db['videos']
                
                # 准备数据
                mongo_data = video_data.copy()
                mongo_data['update_time'] = datetime.now()
                
                if not mongo_data.get('create_time'):
                    mongo_data['create_time'] = datetime.now()
                
                # 更新或插入
                result = videos_collection.update_one(
                    {'video_id': video_data.get('video_id')},
                    {'$set': mongo_data},
                    upsert=True
                )
                
                if result.upserted_id:
                    return str(result.upserted_id)
                else:
                    # 查询更新后的文档ID
                    video = videos_collection.find_one({'video_id': video_data.get('video_id')})
                    return str(video['_id']) if video else False
        except Exception as e:
            logger.error(f"插入/更新视频失败: {e}")
            if self.db_type == 'sqlite':
                self.db.rollback()
            return False

    def check_video_exists(self, video_id):
        """检查视频是否已存在"""
        if self.db is None:
            return False

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                cursor.execute("SELECT id FROM videos WHERE video_id = ?", (video_id,))
                return cursor.fetchone() is not None
            elif self.db_type == 'mongodb':
                videos_collection = self.db['videos']
                return videos_collection.find_one({'video_id': video_id}) is not None
        except Exception as e:
            logger.error(f"检查视频是否存在失败: {e}")
        return False

    def get_video_by_id(self, video_id):
        """根据ID获取视频信息"""
        if self.db is None:
            return None

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                cursor.execute("SELECT * FROM videos WHERE video_id = ?", (video_id,))
                video = cursor.fetchone()
                if video:
                    # 转换为字典
                    columns = [desc[0] for desc in cursor.description]
                    return dict(zip(columns, video))
            elif self.db_type == 'mongodb':
                videos_collection = self.db['videos']
                video = videos_collection.find_one({'video_id': video_id})
                if video:
                    # 转换ObjectId为字符串
                    video['_id'] = str(video['_id'])
                return video
        except Exception as e:
            logger.error(f"获取视频信息失败: {e}")
        return None

    def get_pending_videos(self, limit=100):
        """获取待下载的视频"""
        if self.db is None:
            return []

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                cursor.execute("SELECT * FROM videos WHERE status = 'pending' LIMIT ?", (limit,))
                videos = cursor.fetchall()
                # 转换为字典列表
                columns = [desc[0] for desc in cursor.description]
                return [dict(zip(columns, video)) for video in videos]
            elif self.db_type == 'mongodb':
                videos_collection = self.db['videos']
                videos = list(videos_collection.find({'status': 'pending'}).limit(limit))
                # 转换ObjectId为字符串
                for video in videos:
                    video['_id'] = str(video['_id'])
                return videos
        except Exception as e:
            logger.error(f"获取待下载视频失败: {e}")
        return []

    def update_video_status(self, video_id, status, file_path=None, file_size=None, error_message=None):
        """更新视频状态"""
        if self.db is None:
            return False

        try:
            if self.db_type == 'sqlite':
                cursor = self.db.cursor()
                update_fields = ["status = ?", "update_time = ?"]
                update_values = [status, datetime.now().strftime('%Y-%m-%d %H:%M:%S')]
                
                if file_path:
                    update_fields.append("file_path = ?")
                    update_values.append(file_path)
                if file_size:
                    update_fields.append("file_size = ?")
                    update_values.append(file_size)
                if error_message:
                    # 如果数据库中有error_message字段，则更新它
                    # 检查表结构是否有error_message字段
                    cursor.execute("PRAGMA table_info(videos)")
                    columns = [col[1] for col in cursor.fetchall()]
                    if 'error_message' in columns:
                        update_fields.append("error_message = ?")
                        update_values.append(error_message)
                if status == 'downloaded':
                    update_fields.append("download_time = ?")
                    update_values.append(datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
                
                update_values.append(video_id)
                
                cursor.execute(f"UPDATE videos SET {', '.join(update_fields)} WHERE video_id = ?", update_values)
                self.db.commit()
                return cursor.rowcount > 0
            elif self.db_type == 'mongodb':
                videos_collection = self.db['videos']
                update_data = {
                    '$set': {
                        'status': status,
                        'update_time': datetime.now()
                    }
                }
                
                if file_path:
                    update_data['$set']['file_path'] = file_path
                if file_size:
                    update_data['$set']['file_size'] = file_size
                if error_message:
                    update_data['$set']['error_message'] = error_message
                if status == 'downloaded':
                    update_data['$set']['download_time'] = datetime.now()
                
                result = videos_collection.update_one({'video_id': video_id}, update_data)
                return result.modified_count > 0
        except Exception as e:
            logger.error(f"更新视频状态失败: {e}")
            if self.db_type == 'sqlite':
                self.db.rollback()
        return False

    def close(self):
        """关闭数据库连接"""
        if self.db is not None:
            if self.db_type == 'sqlite':
                self.db.close()
            # MongoDB连接不需要显式关闭
            logger.info("数据库连接已关闭")